#=================================================
#--- for system
import os.path
from datetime import datetime

#--- for data structure
import numpy as np

#--- for file manupilate
from netCDF4 import Dataset

#--- for visualization 
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import cartopy.crs as crs

#--- for WRF output
from wrf import to_np, getvar, smooth2d, get_cartopy, cartopy_xlim, \
                cartopy_ylim, latlon_coords
#=================================================

def gen_ani_map(dflist, pltvar, clmin, clmax, units="", clmap="jet"):
    """
    Create an animation for pltvar in dflist.
    """

    cn_nums = 10
    cn_increment = (clmax - clmin)/cn_nums
    cn_levels = np.arange(clmin, clmax + cn_increment * 0.1, cn_increment)
    
    # Create a figure
    fig = plt.figure(figsize=(12,6))
    ims=[]

    for df in dflist:
        #print(dfname)
        ncfile = Dataset(df)
        prefix = df[-30 : -20]

        yy   = df[-19 : -15]
        mm   = df[-14 : -12]
        dd   = df[-11 : -9]
        hr   = df[-8 : -6]
        mn   = df[-5 : -3]
        sc   = df[-2 : ]
        now_time = yy + "-" + mm +"-" + dd + "_" + hr + ":" + mn + ":" + sc

        # Get the variable to plot
        pltdata = getvar(ncfile, pltvar)

        # Smooth the data since it tends to be noisy near the mountains
        #smooth_data = smooth2d(pltdata, 3, cenweight=4)

        # Get the latitude and longitude points
        lats, lons = latlon_coords(pltdata)

        # Get the cartopy mapping object
        cart_proj = get_cartopy(pltdata)

        # Set the GeoAxes to the projection used by WRF
        ax = plt.axes(projection=cart_proj)

        ax.coastlines('50m', linewidth=0.8)

        # Make the contour outlines and filled contours
        ax.contour(to_np(lons), to_np(lats), to_np(pltdata), colors="black",
                    levels=cn_levels, transform=crs.PlateCarree())
        cntr = ax.contourf(to_np(lons), to_np(lats), to_np(pltdata),
                     levels=cn_levels, transform=crs.PlateCarree(),
                     cmap=clmap, extend="both")

        # Add a color bar
        fig.colorbar(cntr, ax=ax, shrink=.98)

        # Set the map bounds
        ax.set_xlim(cartopy_xlim(pltdata))
        ax.set_ylim(cartopy_ylim(pltdata))

        # Add the gridlines
        ax.gridlines(color="black", linestyle="dotted")

        ax.set_title(pltvar + "(" + units + ") " + now_time, fontsize = 12)

        ims.append([ax])

    # Generate the animation object
    ani = animation.ArtistAnimation(fig, ims, interval=1000, blit=False, \
        repeat=False, repeat_delay=10000)

    # Output the animation object as gif file
    fig_name_pre = pltvar + "_" + prefix
    fig_name_test = fig_name_pre + ".gif"
    if os.path.exists(fig_name_test):
        fig_name_now = fig_name_pre+ "_" + datetime.now().strftime("%Y-%m-%d_%X") + ".gif"
        print("Warning: The animation file " + fig_name_test + "already exists!!!")
        print("         The animation is saved as " + fig_name_now + "!!!")
        ani.save(fig_name_now, writer="imagemagick", fps = 2, dpi = 300)
    else:
        fig_name_now = fig_name_test
        print("---The animation is saved as " + fig_name_now + "!")
        ani.save(fig_name_now, writer="imagemagick", fps = 2, dpi = 300)

    # Show the animation on screen
    plt.show()
